# Direct outputs eligible for public release

# GLOBAL SETTINGS --------------------------------------------------------------

options(
    scipen = 999,
    digits = 16,
    max.print = .Machine$integer.max,
    show.error.locations = TRUE,
    warn = 1
)

RNGkind("L'Ecuyer-CMRG")
seed <- 818675309L
set.seed(seed) # setting main seed

# PACKAGES ---------------------------------------------------------------------
library(data.table)
library(ggplot2)
library(scales)

# PACKAGE SETTINGS -------------------------------------------------------------

# data.table
setDTthreads(threads = 1L)
options(datatable.print.class = TRUE, datatable.print.keys = TRUE)
# so that printing the data.table also shows the variable type on top

# BEGIN FILE -------------------------------------------------------------------

# Just in case something masks 'source()', use extra deliberate base::source().
base::source("~/code/0-utility-functions/graphing-utils.R", local = TRUE)

# Read in and estimation results, and calculate cohort-weighted pre-win means
results_list <-
    readRDS("~/estimation-output/event_study_estimates_moved_tract.rds")

did_coefs <- setDT(results_list[[1]])

did_coefs <-
    did_coefs[
        between(ref_event_time, -7L, 5L, incbounds = TRUE) &
            ref_onset_time == "Cohort-Weighted" &
            model == "reduced_form" &
            rn == "att",
        .(
            ref_event_time,
            estimate,
            cluster_se
        )
    ]

cohort_weights <- setDT(results_list[[1]])
cohort_weights <-
    cohort_weights[
        between(ref_event_time, -7L, 5L, incbounds = TRUE) &
            model == "reduced_form" &
            rn == "catt",
        .(
            ref_onset_time,
            ref_event_time,
            catt_treated_unique_units
        )
    ]

treat_means <-
    rbindlist(
        list(results_list[[3]][[1]], results_list[[3]][[2]]),
        use.names = TRUE
    )

treat_means <-
    treat_means[
        between(ref_event_time, -7L, 5L, incbounds = TRUE) &
            model == "reduced_form",
        .(ref_onset_time, ref_event_time, model, treat_mean_outcome)
    ]
treat_means[, ref_onset_time := as.character(ref_onset_time)]

# Merge in cohort weights to treat_means so as
# to construct cohort-weighted versions of the pre-win means
treat_means <-
    merge(treat_means,
        cohort_weights,
        by = c(
            "ref_onset_time",
            "ref_event_time"
        ),
        all.x = TRUE,
        sort = FALSE
    )

cw_treat_means <-
    treat_means[
        ,
        .(
            ref_onset_time = "Cohort-Weighted",
            cw_treat_mean =
                weighted.mean(
                    x = treat_mean_outcome,
                    w = (
                        catt_treated_unique_units / sum(catt_treated_unique_units) # nolint
                    )
                )
        ),
        by = .(ref_event_time)
    ]

did_coefs <-
    merge(
        did_coefs,
        cw_treat_means,
        by = c("ref_event_time"),
        all.x = TRUE,
        sort = FALSE
    )

# Use the event time representing the most units
did_coefs[
    ,
    pre_win_mean :=
        sum(
            as.integer(
                ref_onset_time == "Cohort-Weighted" & ref_event_time == 0
            ) * cw_treat_mean
        )
]

did_coefs[, c("ref_onset_time", "cw_treat_mean") := NULL]
setnames(did_coefs, "ref_event_time", "event_time")

# Introduce a omitted_event_time row
did_coefs <-
    rbindlist(
        list(
            did_coefs,
            data.table(
                event_time = -2L,
                estimate = 0,
                cluster_se = 0
            )
        ),
        use.names = TRUE,
        fill = TRUE
    )
setorderv(did_coefs, "event_time")

# Housekeeping
results_list <- NULL
cohort_weights <- NULL
treat_means <- NULL
cw_treat_means <- NULL
rm(results_list, cohort_weights, treat_means, cw_treat_means)

# Figure 5.1 -- Moved Tract
figure_5_1 <-
    PlotEventStudyLineTwoY(
        figdata = copy(did_coefs),
        second_y_var = "pre_win_mean",
        event_time_var = "event_time",
        lower_event = -7,
        upper_event = 5,
        omitted_event_time = -2,
        ci_factor = 1.64,
        scale_factor = 100, # to make pp
        base_size = 13,
        xlab_text = "Event Time (ℓ)",
        ylab_text_left = "Event Study Estimate (pp)",
        ylab_text_right = "% Change (relative to pre)",
        y_lower = -0.5,
        y_upper = 4,
        break_width_second_y = 5
    )

ggsave(
    plot = figure_5_1,
    filename = "~/paper/figures/figure-5.1.png",
    width = 6,
    height = 4,
    dpi = 600
)

ggsave(
    plot = figure_5_1,
    filename = "~/paper/figures/figure-5.1.tif",
    width = 6,
    height = 4,
    device = "tiff",
    dpi = 600
)

# Housekeeping
did_coefs <- NULL
figure_5_1 <- NULL
rm(did_coefs, figure_5_1)